
#include "include/sampling.h"
#include "cuda.h"


inline int opt_n_threads(int work_size) {
    const int pow_2 = std::log(static_cast<double>(work_size)) / std::log(2.0);

    return std::max(std::min(1 << pow_2, TOTAL_THREADS), 1);
}



//test function
int gather_points_wrapper_fast(int b, int c, int n, int npoints,
                               at::Tensor points_tensor, at::Tensor idx_tensor, at::Tensor out_tensor){
    const float *points = points_tensor.data<float>();
    const int *idx = idx_tensor.data<int>();
    float *out = out_tensor.data<float>();

    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    gather_points_kernel_launcher_fast(b, c, n, npoints, points, idx, out, stream);
    return 1;
}




void add_cuda(at::Tensor a,torch::Tensor b,torch::Tensor c) {
//    torch::Tensor a= torch::tensor(3.0);
//    torch::Tensor b=torch::tensor(1.0);
//    torch::Tensor c= torch::tensor(0.0);
    const float *aa=a.data<float>();
    const float *bb=b.data<float>();
    float *cc=c.data<float>();

    int x_len=a.size(0),y_len=a.size(1);
    int size[2]={x_len,y_len};
    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    // AddLaunch(aa,bb,cc,size,stream);

}

//c10::ArrayRef<float> arrs = c10::ArrayRef<float>(arr);
//at::Tensor values = at::ones((1600,4),at::kFloat).to("cuda"
//2d array add fuction


void test_the_cuda_add_fun(){

    float arr_1[2][3] = {{3.0, 4.0, 5.0},{1.2, 3.6, 7.8}};
    float arr_2[2][3] = {{1.0, 1.0, 1.0},{2.,1.0,1.}};

//    nc::NdArray<float> array = {{1,2,3},{34,1,2}};
    torch::Tensor tensor_1 = torch::from_blob(arr_1,{2,3}).to("cuda");//rray to tensor
    std::cout<<tensor_1<<std::endl;
    torch::Tensor tensor_2 = torch::from_blob(arr_2,{2,3}).to("cuda");
    std::cout<<tensor_2<<std::endl;
    torch::Tensor  out = torch::zeros({2,3},torch::TensorOptions().dtype(torch::kFloat32)).cuda();

        //get the shape x
    std::cout<<"inital value"<<out<<std::endl;
    add_cuda(tensor_1,tensor_2,out);
    std::cout<<"1+2"<<out<<std::endl;

}



void point_group(){
    //b : batch size
    //c: channels for features
    //n: num of raw points per batch size
    //npoints : num of keypoints
    //features: (bach size, channels, num of raw points)
    int b = 1;
    int c = 32;
    int n = 10;
    int npoints = 5;
    torch::Tensor features = torch::rand({b,c,n},torch::kFloat32).to("cuda");
//    std::cout<<features<<std::endl;
    torch::Tensor idx_tensor = torch::randint(0,n,{b,npoints},torch::kInt32).to("cuda");
//    std::cout<<idx_tensor<<std::endl;
    torch::Tensor out = torch::zeros({b,c,npoints},torch::kFloat32).to("cuda");
    std::cout<<out.tensor_data()<<std::endl;

    gather_points_wrapper_fast(b,  c,  n, npoints,
                               features,  idx_tensor, out);
    std::cout<<out.tensor_data()<<std::endl;
}

void test_FPS_f(){
    //b: batch size;
    //n :of raw points;
    //m: num of keypoints;
    //point_tensor: shape(batch size, num of raw points, channels),
    // the channels for the raw points are 4 ,x,y,z,i respectively,
    //temp tensor:shape(batch size, num of raw points) this is used for save the min distance
    //out: shape( batch size, num of keypoints) this is the idx in raw points for indicing the keypoints from raw points;

}


int furthest_point_sampling_wrapper(int b, int n, int m,
                                    at::Tensor points_tensor, at::Tensor temp_tensor, at::Tensor idx_tensor) {
    //b: batch size;
    //n :of raw points;
    //m: num of keypoints;
    //point_tensor: shape(batch size, num of raw points, channels),
    // the channels for the raw points are 4 ,x,y,z,i respectively,
    //temp tensor:shape(batch size, num of raw points) this is used for save the min distance
    //out: shape( batch size, num of keypoints) this is the idx in raw points for indicing the keypoints from raw points;
    const float *points = points_tensor.data<float>();
    float *temp = temp_tensor.data<float>();
    int *idx = idx_tensor.data<int>();

    cudaStream_t stream = c10::cuda::getCurrentCUDAStream();
    furthest_point_sampling_kernel_launcher(b, n, m, points, temp, idx, stream);
    return 1;
}



int main(){

    test_the_cuda_add_fun();
//    points_test();
//     add_cuda(a,b,c);
//     std::cout<<c<<std::endl;
    return 0;
}

//PYBIND11_MODULE(TORCH_EXTENSION_NAME,m){
//        m.def("add_cuda",&add_cuda,"add_cuda");
//}
